import numpy as np
from scipy.interpolate import make_interp_spline


def smooth_xy(lx, ly):
    """数据平滑处理
    :param lx: x轴数据，数组
    :param ly: y轴数据，数组
    :return: 平滑后的x、y轴数据，数组 [slx, sly]
    """
    x = np.array(lx)
    y = np.array(ly)
    x_smooth = np.linspace(x.min(), x.max(), 300)
    y_smooth = make_interp_spline(x, y)(x_smooth)
    return [x_smooth, y_smooth]


class kalman_filter_2D:
    """二维平面运动的卡尔曼滤波
    : 状态变量 [x, vx, ax, y, vy, ay]
    : 默认测量变量为x和y
    :param delta_t: 采样时间间隔
    :param sigma_a: 加速度标准差
    :param sigma_x: 位置x方向标准差
    :param sigma_y: 位置y方向标准差
    """

    def __init__(self, delta_t=1.0, sigma_a=0.2, sigma_x=1.0, sigma_y=1.0):
        super(kalman_filter_2D, self).__init__()
        # 初始状态变量（参与更新）
        self.x = np.array([[0, 0, 0, 0, 0, 0]]).T
        # 估计不确定性（参与更新）
        self.P = np.array(
            [
                [100, 0, 0, 0, 0, 0],
                [0, 100, 0, 0, 0, 0],
                [0, 0, 100, 0, 0, 0],
                [0, 0, 0, 100, 0, 0],
                [0, 0, 0, 0, 100, 0],
                [0, 0, 0, 0, 0, 100],
            ]
        )
        # 测量矩阵（不参与更新）
        self.H = np.array([[1, 0, 0, 0, 0, 0], [0, 0, 0, 1, 0, 0]])
        # 状态转移矩阵（根据系统方程确定）
        self.F = self.make_F()
        # 过程噪声不确定性（根据系统方程和超参数 sigma_a 确定）
        self.Q = self.make_Q(delta_t, sigma_a)
        # 测量不确定性（根据超参数 sigma_x 和 sigma_y 确定）
        self.R = self.make_R(sigma_x, sigma_y)

    def make_F(self, delta_t):
        """状态转移矩阵
        : 由系统方程确定
        :param delta_t
        :return F
        """
        return np.array(
            [
                [1, delta_t, 0.5 * delta_t**2.0, 0, 0, 0],
                [0, 1, delta_t, 0, 0, 0],
                [0, 0, 1, 0, 0, 0],
                [0, 0, 0, 1, delta_t, 0.5 * delta_t**2.0],
                [0, 0, 0, 0, 1, delta_t],
                [0, 0, 0, 0, 0, 1],
            ]
        )

    def make_Q(self, delta_t, sigma_a):
        """过程噪声不确定性
        : 根据系统方程和超参数sigma_a确定
        :param delta_t, sigma_a
        :return Q
        """
        return sigma_a**2 * np.array(
            [
                [delta_t**4 / 4, delta_t**3 / 2, delta_t**2 / 2, 0, 0, 0],
                [delta_t**3 / 2, delta_t**2, delta_t, 0, 0, 0],
                [delta_t**2 / 2, delta_t, 1, 0, 0, 0],
                [0, 0, 0, delta_t**4 / 4, delta_t**3 / 2, delta_t**2 / 2],
                [0, 0, 0, delta_t**3 / 2, delta_t**2, delta_t],
                [0, 0, 0, delta_t**2 / 2, delta_t, 1],
            ]
        )

    def make_R(self, sigma_x, sigma_y):
        """测量不确定性
        : 根据超参数sigma_x和sigma_y确定
        :param sigma_x, sigma_y
        :return R
        """
        return np.array([[sigma_x**2, 0], [0, sigma_y**2]])

    def predict(self, x, P, F, Q):
        """预测函数
        :param x, P, F, Q
        :return x_, P_
        """
        x_ = F @ x
        P_ = F @ P @ F.T + Q
        return x_, P_

    def measure(self, z_x: iter, z_y: iter):
        """测量函数
        :param z_x, z_y
        :return z
        """
        return np.array([[next(z_x), next(z_y)]]).T

    def update(self, x, z, P, H, R):
        """更新函数
        :param x, z, P, H, R
        :return K, x, P
        """
        K = P @ H.T @ np.linalg.pinv(H @ P @ H.T + R)
        x = x + K @ (z - H @ x)
        P = (np.eye(len(P)) - K @ H) @ P @ (np.eye(len(P)) - K @ H).T + K @ R @ K.T
        return K, x, P

    def iterate(self, num, z_x, z_y):
        """迭代计算函数
        :param num: 迭代次数
        :param z_x, z_y: 测量x和y
        :return pred_x, pred_y: 预测x和y
        """
        pred_x = []
        pred_y = []
        x_, P_ = self.predict(self.x, self.P, self.F, self.Q)
        for _ in range(num):
            # measure
            z = self.measure(z_x, z_y)
            # update
            K, x, P = self.update(x_, z, P_, self.H, self.R)
            pred_x.append(x[0][0])
            pred_y.append(x[3][0])
            # predict
            x_, P_ = self.predict(x, P, self.F, self.Q)

        return pred_x, pred_y
